#ifndef __HUFFMAN_H__
#define __HUFFMAN_H__

#include "Bitstream.h"

#include <map>

//---------------------------------------------------------------------------
//---------------------------------------------------------------------------
//---------------------------------------------------------------------------

struct HuffmanNode
{
  HuffmanNode() : LeftChild(-1),RightChild(-1) {}
  HuffmanNode(int value,int leftChild,int rightChild) :
    Value(value),LeftChild(leftChild),RightChild(rightChild) {}
  
  //----------------------------------------------
  
  int Value;
  unsigned int Code;
  int CodeSize;
  int LeftChild,RightChild;
};

//---------------------------------------------------------------------------
//---------------------------------------------------------------------------
//---------------------------------------------------------------------------

typedef std::multimap<int,int> MMap;
typedef std::vector<HuffmanNode> TreeVector;
typedef std::map<int,std::pair<int,int> > CMap;

//---------------------------------------------------------------------------
//---------------------------------------------------------------------------
//---------------------------------------------------------------------------


//===========================================================
//===================PACKER==================================
//===========================================================

//---------------------------------------------------------------------------
//---------------------------------------------------------------------------
//---------------------------------------------------------------------------

class HuffmanPacker
{
 public:
  //----------------------------------------------
  //----------------------------------------------
  //----------------------------------------------

  void Pack(const unsigned char * src,int srcSize);
  void GetBuffer(void** buffer,int& bufferSize) const;

  //----------------------------------------------
  //----------------------------------------------
  //----------------------------------------------
 protected:
  //----------------------------------------------
  //----------------------------------------------
  //----------------------------------------------

  void CollectData(const unsigned char * src,int srcSize);
  void UpdateTreeCode(int index,int treeDepth,unsigned int code);
  void BuildTree(void);
  void WriteTree(int index);
  void WriteHeader(void);
  void WriteStream(const unsigned char * src,int srcSize);

  //----------------------------------------------
  //----------------------------------------------
  //----------------------------------------------

  int SrcSize;
  int InputBits;
  MMap Collector;
  TreeVector Tree;//last element is the root of the tree
  OutputBitStream Stream;

  //----------------------------------------------
  //----------------------------------------------
  //----------------------------------------------
};

//-------------------------------------------------------------------------
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------

void HuffmanPacker::CollectData(const unsigned char * src,int srcSize)
{
  std::map<int,int> tempCollector;

  for(int i=0; i<srcSize; ++i)//collects all unique data and their number of occurences
    {
      std::pair< std::map<int,int>::iterator , bool > p = tempCollector.insert( std::make_pair(src[i],1) );
      if(!p.second)//existed already
	++p.first->second;//one more found
    }

  SrcSize=srcSize;

  //------------------------

  //inverses key and value for sorting along the occurence counter
  //and prepare elements for the Tree
  for(std::map<int,int>::iterator it=tempCollector.begin(); it!=tempCollector.end(); ++it)
    {
      HuffmanNode Node = HuffmanNode(it->first,-1,-1);//value,leftChild,rightChild
      Collector.insert( std::make_pair( it->second , (int)Tree.size()) );//Tree.size=index in tree
      Tree.push_back( Node );
    }
}

//-------------------------------------------------------------------------
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------

void HuffmanPacker::UpdateTreeCode(int index,int treeDepth,unsigned int code)
{
  //bits in the roots of the tree will be low-level bits in the code
  //thus first bits in the stream => ok for read
  if(Tree[index].LeftChild==-1)//leaf
    {
      Tree[index].Code=code;
      Tree[index].CodeSize=treeDepth;
    }
  else//not a leaf => recurse
    {
      UpdateTreeCode(Tree[index].LeftChild,treeDepth+1,code);
      UpdateTreeCode(Tree[index].RightChild,treeDepth+1, code | (1<<treeDepth) );
    }
}

//-------------------------------------------------------------------------
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------

void HuffmanPacker::BuildTree(void)
{
  //remove nodes from collector (by building tree) until there is only one left
  while(Collector.size()>1)
    {
      //find two elements with the lowest counts
      MMap::iterator lowest=Collector.begin();
      MMap::iterator lowest2=lowest;
      ++lowest2;//second lowest
      int lowestIndex=lowest->second;//index in the tree vector
      int lowest2Index=lowest2->second;
      int count=lowest->first+lowest2->first;

      //------------------------

      //create a node as sum of the two smallest nodes
      //inserts node into collector and tree
      Tree.push_back(HuffmanNode(0,lowestIndex,lowest2Index));

      Collector.insert(std::make_pair(count,(int)Tree.size()-1));//tree.size-1=index in tree

      //remove two old elements from collector
      Collector.erase(lowest);
      Collector.erase(lowest2);
    }
  //update nodes levels from root and calculate codes
  UpdateTreeCode((int)Tree.size()-1,0,0);
}

//-------------------------------------------------------------------------
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------

void HuffmanPacker::WriteTree(int index)
{
  if(Tree[index].LeftChild==-1)
    {
      Stream.PushBit(1);//it's a leaf
      Stream.PushValue(Tree[index].Value,8);//save value
    }
  else
    {
      Stream.PushBit(0);//it's not a leaf
      WriteTree(Tree[index].LeftChild);//recursion
      WriteTree(Tree[index].RightChild);
    }
}

//-------------------------------------------------------------------------
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------

void HuffmanPacker::WriteHeader(void)
{
  WriteExpGolomb(SrcSize,Stream);
  //recursively write tree
  WriteTree((int)Tree.size()-1);
}

//-------------------------------------------------------------------------
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------

void HuffmanPacker::WriteStream(const unsigned char * src,int srcSize)
{
  CMap codeMap;
  //prepare the map for fast finding the code of input data
  for(TreeVector::iterator it=Tree.begin(); it!=Tree.end(); ++it)
    {
      if(it->LeftChild!=-1)//if it's not a leaf => forget it
	continue;
      codeMap.insert(std::make_pair(it->Value,std::make_pair(it->Code,it->CodeSize)));
    }

  //int value=0;
  for(int i=0;i<srcSize;++i)
    {
      CMap::iterator it=codeMap.find(src[i]);//fast thanks to the map
      Stream.PushValue(it->second.first,it->second.second);//code and codesize
    }
}

//-------------------------------------------------------------------------
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------

void HuffmanPacker::Pack(const unsigned char * src,int srcSize)
{
  CollectData(src,srcSize);

  BuildTree();

  Stream.Begin();

  WriteHeader();

  //int headerSize=(int)Stream.GetDataSize();

  WriteStream(src,srcSize);

  Stream.End();

  //int dataSize=(int)Stream.GetDataSize();
}

//-------------------------------------------------------------------------
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------

void HuffmanPacker::GetBuffer(void **buffer,int &bufferSize) const
{
  *buffer=(void *)Stream.GetData();
  bufferSize=(int)Stream.GetDataSize();
}

//-------------------------------------------------------------------------
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------

//===========================================================
//======================DEPACKER=============================
//===========================================================

//---------------------------------------------------------------------------
//---------------------------------------------------------------------------
//---------------------------------------------------------------------------

class HuffmanDepacker
{
  //----------------------------------------------
  //----------------------------------------------
  //----------------------------------------------
 public:
  //----------------------------------------------
  //----------------------------------------------
  //----------------------------------------------

  void SetBuffer(const void *buffer);
  //Depack Size is the number of elements that have been compressed
  int GetDepackSize(void) const { return DstSize; };
  void Depack(unsigned char *destination);

  //----------------------------------------------
  //----------------------------------------------
  //----------------------------------------------
 protected:
  //----------------------------------------------
  //----------------------------------------------
  //----------------------------------------------

  void ReadTree(int index,int treeDepth,unsigned int code);
  unsigned char	ReadCode(int index);

  //----------------------------------------------
  //----------------------------------------------
  //----------------------------------------------

  unsigned int DstSize;
  unsigned int OutputBits;
  TreeVector Tree;
  InputBitStream Stream;

  //----------------------------------------------
  //----------------------------------------------
  //----------------------------------------------
};

//-------------------------------------------------------------------------
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------

void HuffmanDepacker::SetBuffer(const void *buffer)
{
  Stream.Begin((unsigned short *)buffer);
  ReadExpGolomb(Stream,DstSize);
}

//-------------------------------------------------------------------------
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------

void HuffmanDepacker::ReadTree(int index,int treeDepth,unsigned int code)
{
  if(Stream.Pop(1))//leaf
    {
      Tree[index].Value=Stream.Pop(8);//read value
      Tree[index].Code=code;//keep code
    }
  else//not a leaf => recurse
    {
      Tree[index].LeftChild=(int)Tree.size();
      Tree.push_back(HuffmanNode());
      ReadTree(Tree[index].LeftChild,treeDepth+1,code);

      Tree[index].RightChild=(int)Tree.size();
      Tree.push_back(HuffmanNode());
      ReadTree(Tree[index].RightChild,treeDepth+1,code | (1<<treeDepth));
    }
}

//-------------------------------------------------------------------------
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------

unsigned char HuffmanDepacker::ReadCode(int index)
{
  if(Tree[index].LeftChild==-1)//leaf
    return Tree[index].Value;
  else//recurse
    {
      if(!Stream.Pop(1))
	return ReadCode(Tree[index].LeftChild);
      else
	return ReadCode(Tree[index].RightChild);
    }
}

//-------------------------------------------------------------------------
//-------------------------------------------------------------------------
//-------------------------------------------------------------------------

void HuffmanDepacker::Depack(unsigned char * destination)
{
  Tree.push_back(HuffmanNode());//create root of the tree
  ReadTree(0,0,0);

  for(unsigned int i=0;i<DstSize;++i)
    destination[i]=ReadCode(0);
}

//--------------------------------------------------------------------------
//--------------------------------------------------------------------------
//--------------------------------------------------------------------------

#endif //__Huffman_H__
